# Title     : TODO
# Objective : TODO
# Created by: Administrator
# Created on: 2019/6/18

library(plyr)
library(gridExtra)
library(ggpubr)
library(egg)
library(lemon)
library(purrr)
library(ComplexHeatmap)
library(Hmisc)
library(optparse)
library(tidyverse)

createWhenNoExist <- function(f){
    ! dir.exists(f) && dir.create(f)
}

options(digits = 3)

option_list <- list(
make_option("--i", default = "AllMet.csv", type = "character", help = "metabolite data file"),
make_option("--g", default = "SampleInfo.csv", type = "character", help = "sample group file"),
make_option("--sc", default = "sample_color.txt", type = "character", help = "sample color file"),
make_option("--mc", default = "meta_color.txt", type = "character", help = "metabolite color file")
)
opt <- parse_args(OptionParser(option_list = option_list))

sampleInfo <- read.csv(opt$g, header = T, stringsAsFactors = F) %>%
    select(c("SampleID", "ClassNote"))

classNotes <- sampleInfo %>%
.$ClassNote %>%
unique()
cn <- combn(classNotes, 2)
for (i in 1 : ncol(cn)) {
    row <- cn[, i]
    group1Name <- row[1]
    group2Name <- row[2]
    parent <- paste0(group2Name, "_", group1Name)
    createWhenNoExist(parent)
    fileName <- paste0(parent, "/diff_z-score_", group2Name, "_", group1Name, ".txt")
    scoreData <- read.table(fileName, header = T, sep = "\t", stringsAsFactors = F, comment.char = "", check.names = F,
    row.names = 1)

    metaboliteInfo <- read.csv(opt$i, header = T, stringsAsFactors = F) %>%
    select(c("Metabolite", "Class"))
    classScoreData <- data.frame(Name = colnames(scoreData), stringsAsFactors = F) %>%
    left_join(metaboliteInfo, by = c("Metabolite"))
    classNotes <- scoreData %>%
        rownames_to_column("SampleID") %>%
        left_join(sampleInfo, by = c("SampleID")) %>%
        .$ClassNote
    print(head(classScoreData))

    pdfFileName <- paste0(parent, "/Diff_Meta_Zscore_", group2Name, "_", group1Name, "_Cluster.pdf")
    pdf(pdfFileName, width = 8, height = 10)
    plotData <- scoreData %>%
    t()
    colors <- colorRampPalette(c("green", "black", "red"), space = "rgb")(256)

    sampleColDf <- read.table(opt$sc, header = T, sep = "\t", stringsAsFactors = F, comment.char = "") %>%
    select(c("ClassNote", "col"))
    sampleCols <- sampleColDf %>%
    deframe()

    colHa <- HeatmapAnnotation(ClassNote = classNotes, col = list(ClassNote = sampleCols),
    annotation_legend_param = list(ClassNote = list(nrow = 1, title = "",
    labels_gp = gpar(fontsize = 12, fontfamily = "Times"))), show_annotation_name = F
    )

    metaColDf <- read_tsv(opt$mc) %>%
    select(c("Class", "col"))
    metaCols <- metaColDf %>%
    deframe()

    rowHa <- rowAnnotation(Class = classScoreData$Class, col = list(Class = metaCols),
    annotation_legend_param = list(Class = list(title = "", labels_gp = gpar(fontsize = 12, fontfamily = "Times"))),
    annotation_name_gp = gpar(fontsize = 6, fontfamily = "Times"), show_annotation_name = F
    )

    htList <- Heatmap(plotData, col = colors, show_column_names = F, cluster_rows = T, cluster_columns = F,
    name = "z-score", top_annotation = colHa, right_annotation = rowHa,
    row_names_gp = gpar(fontsize = 6, fontfamily = "Times"),
    heatmap_legend_param = list(title_gp = gpar(fontsize = 12, fontfamily = "Times"),
    labels_gp = gpar(fontsize = 12, fontfamily = "Times"))
    )
    draw(htList, heatmap_legend_side = "left", annotation_legend_side = "top",
    padding = unit(c(0.5, 0.5, 0.5, 0.5), "cm")
    )
    dev.off()
}



